import torch
import torch.nn as nn
import torch.nn.functional as F
import os.path as osp
from .ResNet import ResNet101
from .ASPP import ASPP
from .Decoder import Decoder
from torch.hub import load_state_dict_from_url

class DeepLab_v3(nn.Module):

    # pretrained_model = osp.expanduser('~//data//models//pytorch//deeplab-resnet.pth')
    weight_url='https://class-1275-42261.obs.cn-north-1.myhuaweicloud.com:443/Lab-2210/modelarts22905308/train_output/V0015/ckpt_45.pth?AWSAccessKeyId=LKSWOF8QARN7XM2GZO7W&Expires=1594973099&response-content-disposition=inline&x-amz-security-token=gQpjbi1ub3J0aC00jKlokKTkaTRVGhWka4sjeB8VeIQIVOzuBp0FDqolGiYjZ7m81y4kIpnKENr_2SrOMdAkqre4Hv9IW3Hs31FRoWkzcOJ2qy1lO7DFzCN-R7GlopMsvg8dwcE1jqDR7f7e2sKtaS7ZTy5Dfc8UJf5YgAjn7O36eDg2-Ncj2IcvBZAsCh4Q_EMwvodCoNUumyRhUPDLxH9W3dfUBi-os8wO3VH7nJzNyBRksxsNjIvXTGbg36RV6soDjsDUADDgiWxvHgMbZ7CgrvBicl-6pRGiICDrMeaB-zdqhlGQEF4-Z_M0Ty0qgyIR7up_f3JDQBzvBk_UdZQg5Kbc7e5IdJSI5Y2bkGMrpgBZzg_lMOpDhBdQH9GFk8myt6HVDl1we19dqEI1eypW0QAFNGwPhldz4gTV6qyAzXbKLgjz11SRfx7X6ZdgHeqWTKdtc058jKBIuNAli47BIkR84ObEu_z1_N1LYaNvbdCwC0SK01KQqP0lfANmgDeHyyF6p_ptHFLmQZ3MM-xN2Twgnb8cjeBNA3MuMKLVz3DDX1S3OAetreB2ByMpzYaC1j-mg6NIM-aSXDrZR731VwkRVoVoh2ikSxPIr4T1IOl9t3w6xVax5L557GD0WesmZcoqIAdsAGTXUZrmYY5HzewxvroZ1SnEF0rHiygTLMicjwZIDjh4sBT21qgrSMVkmVstoZr7--3zkX3Zx2-GmB0tTMSnOXiijWOo5IUIqeH6TCxCKvbIodmQVaw5BTpuyT6D4-YOLqbNM-x_eKwmw2RE0M9UixfNx8zf6H0oDf4Sd8DTez59ktE0&Signature=57cEEoaiE8ggjT%2Bj8uSH0NY7srw%3D'
    def __init__(self, backbone='resnet', output_stride=16, num_classes=8, freeze_bn=False,pretrain=False):
        super(DeepLab_v3, self).__init__()
        # batch_norm = nn.BatchNorm2d
        self.backbone = ResNet101(output_stride)
        self.aspp = ASPP(output_stride)
        self.decoder = Decoder(num_classes)
        self.freeze_bn = freeze_bn
        if pretrain:
            # state_dict=torch.load(self.pretrained_model)['state_dict']
            pretrain_dict = load_state_dict_from_url(self.weight_url,
                                                     progress=True)
            self.load_state_dict(pretrain_dict)
    def forward(self,input):
        x, low_level_feat = self.backbone(input)
        x = self.aspp(x)
        x = self.decoder(x, low_level_feat)
        x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)

        return x